import paramiko
import os
import logging
from scp import SCPClient
import socket
import time

logger = logging.getLogger(__name__)

class SSHUtils:
    def __init__(self):
        self.ssh = paramiko.SSHClient()
        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.connected = False
    
    def connect(self, host, port, username, password=None, key_path=None, timeout=10):
        """连接到SSH服务器"""
        try:
            logger.info(f"尝试SSH连接到 {username}@{host}:{port}")
            
            # 检查必要参数
            if not all([host, port, username]):
                return False, "缺少必要的连接参数"
            
            # 清理参数
            host = host.strip()
            username = username.strip()
            port = int(port)
            
            # 检查主机是否可达
            if not self._check_host_reachable(host, port, timeout):
                return False, f"主机 {host}:{port} 不可达"
            
            connection_params = {
                'hostname': host,
                'port': port,
                'username': username,
                'timeout': timeout
            }
            
            if key_path and os.path.exists(key_path):
                logger.info(f"使用密钥认证: {key_path}")
                try:
                    # 尝试不同的密钥格式
                    try:
                        private_key = paramiko.RSAKey.from_private_key_file(key_path)
                    except paramiko.ssh_exception.PasswordRequiredException:
                        return False, "SSH密钥需要密码"
                    except:
                        # 尝试其他密钥类型
                        private_key = paramiko.Ed25519Key.from_private_key_file(key_path)
                    
                    connection_params['pkey'] = private_key
                except Exception as key_error:
                    logger.error(f"密钥加载失败: {key_error}")
                    return False, f"SSH密钥加载失败: {key_error}"
            
            elif password:
                logger.info("使用密码认证")
                connection_params['password'] = password
            else:
                return False, "未提供认证方式（密码或密钥）"
            
            # 建立连接
            self.ssh.connect(**connection_params)
            self.connected = True
            
            # 测试连接是否有效
            stdin, stdout, stderr = self.ssh.exec_command('echo "SSH连接测试成功"', timeout=5)
            exit_status = stdout.channel.recv_exit_status()
            
            if exit_status == 0:
                logger.info(f"SSH连接成功: {username}@{host}:{port}")
                return True, "连接成功"
            else:
                error_output = stderr.read().decode().strip()
                return False, f"连接测试失败: {error_output}"
                
        except paramiko.AuthenticationException:
            error_msg = "认证失败：用户名或密码错误"
            logger.error(error_msg)
            return False, error_msg
        except paramiko.SSHException as e:
            error_msg = f"SSH连接错误: {e}"
            logger.error(error_msg)
            return False, error_msg
        except socket.timeout:
            error_msg = f"连接超时：{host}:{port}"
            logger.error(error_msg)
            return False, error_msg
        except socket.gaierror:
            error_msg = f"主机名解析失败: {host}"
            logger.error(error_msg)
            return False, error_msg
        except Exception as e:
            error_msg = f"连接失败: {e}"
            logger.error(error_msg)
            return False, error_msg
    
    def _check_host_reachable(self, host, port, timeout=5):
        """检查主机是否可达"""
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(timeout)
            result = sock.connect_ex((host, port))
            sock.close()
            return result == 0
        except:
            return False
    
    def test_connection(self, host, port, username, password=None, key_path=None):
        """测试连接并返回详细结果"""
        start_time = time.time()
        success, message = self.connect(host, port, username, password, key_path)
        end_time = time.time()
        
        result = {
            'success': success,
            'message': message,
            'response_time': round((end_time - start_time) * 1000, 2),  # 毫秒
            'timestamp': time.strftime('%Y-%m-%d %H:%M:%S')
        }
        
        self.close()
        return result
    
    def execute_command(self, command, timeout=10):
        """执行远程命令"""
        if not self.connected:
            return None, None, "未建立连接"
        
        try:
            stdin, stdout, stderr = self.ssh.exec_command(command, timeout=timeout)
            output = stdout.read().decode().strip()
            error = stderr.read().decode().strip()
            return output, error, None
        except Exception as e:
            return None, None, str(e)
    
    def close(self):
        """关闭连接"""
        try:
            if hasattr(self, 'ssh') and self.connected:
                self.ssh.close()
                self.connected = False
        except:
            pass
    def transfer_file(self, local_path, remote_path, progress_callback=None):
        """传输文件到远程服务器，支持进度回调"""
        if not self.connected:
            return False, "未建立SSH连接"
        
        try:
            # 获取本地文件大小用于进度计算
            if not os.path.exists(local_path):
                return False, f"本地文件不存在: {local_path}"
            
            file_size = os.path.getsize(local_path)
            logger.info(f"开始传输文件: {local_path} -> {remote_path} (大小: {file_size} bytes)")
            
            # 自定义SCP进度处理
            def scp_progress(filename, size, sent):
                if progress_callback and size > 0:
                    progress_callback(sent, size)
            
            # 使用SCP传输文件
            with SCPClient(self.ssh.get_transport(), progress=scp_progress) as scp:
                scp.put(local_path, remote_path)
            
            logger.info(f"文件传输完成: {local_path} -> {remote_path}")
            return True, None
            
        except Exception as e:
            error_msg = f"文件传输失败: {str(e)}"
            logger.error(error_msg)
            return False, error_msg
# 创建工具实例
ssh_utils = SSHUtils()
